import cv2
import numpy as np
from PIL import Image


def show(img, name=""):
    cv2.imshow(name, img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


img = cv2.imread("rsc/abcd.jpg")
# 图片好大，缩小点
w, h = img.shape[:2]
img = cv2.resize(img, (int(w*0.4), int(h*0.4)), interpolation=cv2.INTER_AREA)

# 灰度图
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
show(gray, "gray")

# 边缘检测
edged = cv2.Canny(img, 200, 300)
show(edged, "edged")

edged = cv2.morphologyEx(edged, cv2.MORPH_CLOSE, np.ones((9, 3)), iterations=2)
show(edged, "edged_close")

# 找到矩形轮廓
screen = None
contours = cv2.findContours(edged, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)[0]
img_cp = img.copy()
cv2.drawContours(img_cp, contours, -1, (0, 255, 0), 3)
show(img_cp, "all_cnt")
# 找大的先，所以排序一下
contours = sorted(contours, key=cv2.contourArea, reverse=True)
for contour in contours:
    # 计算封闭周长用于画框的阈值
    length = cv2.arcLength(contour, True)
    approx = cv2.approxPolyDP(contour, 0.02 * length, True)
    if len(approx) == 4:
        screen = approx
        break

img_cp = img.copy()
cv2.drawContours(img_cp, [screen], -1, (0, 255, 0), 3)
show(img_cp, "scan_place")

# 计算一下长宽
# print(screen)
screen = screen.reshape((4, 2))
print(screen)
# # 先按左上 右上 右下 左下的顺序排列
# xy_sum = np.sum(screen, axis=1)
# tl = screen[np.argmax(xy_sum)]
# br = screen[np.argmin(xy_sum)]
# xy_diff = np.diff(screen, axis=1)
# tr = screen[np.argmax(xy_diff)]
# bl = screen[np.argmin(xy_diff)]
# print(tl, tr, br, bl)

# 左上 右上 右下 左下
screen_list = sorted(screen, key=lambda x: (x[1], x[0]))
# screen_list[2], screen_list[3] = screen_list[3], screen_list[2]
screen_sorted = np.array(screen_list, dtype=np.float32)
# 后两个换下位置
screen_sorted[[2, 3], :] = screen_sorted[[3, 2], :]
print(screen_sorted)
tl, tr, br, bl = screen_sorted

# 这里是因为不熟悉numpy的操作搞得, 而且也不知道取到的tl, tr, br, bl是引用! 浪费了很多时间
# print(tl, tr, br, bl)
# print(screen_sorted)
# screen_sorted[2] = br
# screen_sorted[3] = bl
# print(screen_sorted)
# print(tl, tr, br, bl)

width1 = np.sqrt((tr[0] - tl[0]) ** 2 + (tr[1] - tl[1]) ** 2)
width2 = np.sqrt((br[0] - bl[0]) ** 2 + (br[1] - bl[1]) ** 2)
height1 = np.sqrt((br[0] - tr[0]) ** 2 + (br[1] - tr[1]) ** 2)
height2 = np.sqrt((bl[0] - tl[0]) ** 2 + (bl[1] - tl[1]) ** 2)

height = int(max(height1, height2))
width = int(max(width1, width2))
print(height, width)

# 变换后的位置
dst = np.array(
    [
        [0, 0],
        [width, 0],
        [width, height],
        [0, height]
    ],
    dtype=np.float32
)

# 获取变换矩阵
M = cv2.getPerspectiveTransform(screen_sorted, dst)
# 变换
scan = cv2.warpPerspective(img, M, (width, height))
show(scan, "scan")
print(scan.shape)

scan = Image.fromarray(scan)
# 旋转一下, 用PIL比较方便
scan = scan.rotate(-90).resize((scan.height, scan.width))
scan = np.array(scan)
print(scan.shape)
show(scan, "scan_rotated")
cv2.imwrite("rsc/abcd_scan.jpg", scan)
